import tensorflow as tf
import numpy as np

model_save_path = "checkpoint/rnn_onehot_1prel.ckpt"

model = tf.keras.models.Sequential([
    tf.keras.layers.SimpleRNN(3),
    tf.keras.layers.Dense(5, activation='softmax')
])

model.load_weights(model_save_path)

input_word = 'abcde'
id = {'a': 0, 'b': 1, 'c': 2, 'd': 3, 'e': 4}
onehot = {0: [1, 0, 0, 0, 0], 1: [0, 1, 0, 0, 0], 2: [0, 0, 1, 0, 0], 3: [0, 0, 0, 1, 0], 4: [0, 0, 0, 0, 1]}

# 预测
while (1):
    alphabet1 = input('the test alphabet:')
    alphabet = [onehot[id[alphabet1]]]
    # 使alphabet符合simpleRNN的输入要求：送入一个样本，样本数为1；输入一个字母出预测，时间展开步为1；
    # 表示独热码有5个输入特征，每个时间输入特征个数为5
    alphabet = np.reshape(alphabet, (1, 1, 5))
    result = model.predict([alphabet])
    pred = tf.argmax(result, axis=1)
    pred = int(pred)
    tf.print(alphabet1 + '->' + input_word[pred])
    print('\n')
